#include "simulation.h"
#include "factor.h"
#include "boost/random.hpp"
#include "boost/random/normal_distribution.hpp"
#include "Eigen/Geometry"
#include "g2o/stuff/sampler.h"
#include <g2o/core/block_solver.h>
#include <g2o/core/optimization_algorithm_levenberg.h>
#include <g2o/core/optimization_algorithm_gauss_newton.h>
#include <g2o/solvers/dense/linear_solver_dense.h>
#include <g2o/solvers/cholmod/linear_solver_cholmod.h>

Simulation::Simulation() {}

bool Simulation::optimize()
{
    Eigen::Matrix3d K = Eigen::Matrix3d::Identity();
    Eigen::Matrix2d featInfo = featVar.inverse();
    g2o::SparseOptimizer optimizer;
    typedef g2o::BlockSolver_6_3 Block_6_3;
    Block_6_3::LinearSolverType* linearSolver_6_3 = new g2o::LinearSolverCholmod<Block_6_3::PoseMatrixType>();
    Block_6_3 * solver_ptr = new Block_6_3(linearSolver_6_3);
    g2o::OptimizationAlgorithmLevenberg* solver = new g2o::OptimizationAlgorithmLevenberg(solver_ptr);
    optimizer.setAlgorithm(solver);
    std::vector<Eigen::Vector3d> points;
    Pose &pose1 = poses[0];
    Pose &pose2 = poses[1];
    Sophus::SE3d T12 = pose1.pose * pose2.pose.inverse();
    int vertexId = 0;
    for(size_t i = 0; i < pose1.noiseFeats.size(); ++i) {
        Eigen::Vector2i &ft1 = pose1.noiseFeats[i];
        Eigen::Vector2i &ft2 = pose2.noiseFeats[i];
        Eigen::Vector3d p1(double(ft1(0)), double(ft1(1)), 1.0);
        Eigen::Vector3d p2(double(ft2(0)), double(ft2(1)), 1.0);
        Eigen::Vector3d p2r = p1.cross(T12.so3() * p2);
        Eigen::Vector3d p2t = -p1.cross(T12.translation());
        int num = 0;
        double sum = 0.0;
        for(size_t i = 0; i < 3; ++i) {
            double s = p2t[i] / p2r[i];
            if(s > 0.0001 && s < 100) {
                sum += s;
                num++;
            }
        }

        double s2 = sum / double(num);
        points.push_back((pose2.pose.inverse() * p2) * s2);
        auto *vertex = new point(points.back());
        vertex->setId(vertexId++);
        optimizer.addVertex(vertex);
    }

    size_t edgeId = 0;
    size_t poseId = vertexId;
    for(size_t i = 0; i < poses.size(); ++i) {
        Pose &pose = poses[i];
        node * pose_node = new node(pose.pose);
        pose_node->setId(vertexId++);
        if (vertexId - 1 == 0)
            pose_node->setFixed(true);

        optimizer.addVertex(pose_node);
        for(size_t k = 0; k < pose.noiseFeats.size(); ++k) {
            Eigen::Vector2d px(pose.noiseFeats[k](0), pose.noiseFeats[k](1));
                    BAFactor *factor = new BAFactor(px, featInfo, K);
            factor->setId(edgeId++);
            factor->setVertex(0, pose_node);
            factor->setVertex(1, optimizer.vertices()[k]);
            optimizer.addEdge(factor);
        }
    }


    size_t nodeid = 0;
    Eigen::Matrix<double, 6, 6> gpsInfo = gpsVar.inverse();
    for(size_t i = 0; i < gpsMeasures.size(); ++i) {
        double time = gpsMeasures[i].timeStamp;
        while(nodeid < poses.size() && time < poses[nodeid].timeStamp) nodeid++;
        double alpha = (time - poses[nodeid - 1].timeStamp)
                / (poses[nodeid].timeStamp - poses[nodeid - 1].timeStamp);
        if(alpha < 0.00001) {
            unary_gps_factor *gpsfactor = new unary_gps_factor(gpsMeasures[i].pose, gpsInfo);
            gpsfactor->setId(edgeId++);
            gpsfactor->setVertex(0, optimizer.vertices()[poseId + nodeid - 1]);
            optimizer.addEdge(gpsfactor);
        }

        else if(alpha > 0.99999) {
            unary_gps_factor *gpsfactor = new unary_gps_factor(gpsMeasures[i].pose, gpsInfo);
            gpsfactor->setId(edgeId++);
            gpsfactor->setVertex(0, optimizer.vertices()[poseId + nodeid]);
            optimizer.addEdge(gpsfactor);
        }
        else {
            binary_gps_factor *gpsfactor = new binary_gps_factor(gpsMeasures[i].pose, gpsInfo, alpha);
            gpsfactor->setId(edgeId++);
            gpsfactor->setVertex(0, optimizer.vertices()[poseId + nodeid - 1]);
            gpsfactor->setVertex(1, optimizer.vertices()[poseId + nodeid]);
            optimizer.addEdge(gpsfactor);
        }

    }


    std::cout << "begin optimize() \n";
    optimizer.setVerbose(true);
    optimizer.initializeOptimization();
    std::cout << "initial over!\n";
    optimizer.optimize(30);
    return true;

}

double get_gpstimestamp(double start, double end) {
    boost::mt19937 rng(clock());
    boost::uniform_real<> ur(start, end);
    boost::variate_generator<boost::mt19937, boost::uniform_real<>> rand_(rng, ur);
    return rand_();
}

void Simulation::run(Eigen::Matrix2d &featVar)
{
    this->featVar = featVar;
    double r = 10.0;
    double b1[3] = {1.5, -0.3, 2.5}, b2[3] = {3.7, 0, 3}, b3[3] = {4.3, 2.4, 6}, b4[3] = {1.4, 1.5, 4},
            b5[3] = {-0.5, 1, 5}, b6[3] = {- 1, 0, 4.2}, b7[3] = {0, -1, 3.4}, b8[3] = {-3.3, 2.6};
    Eigen::Map<Eigen::Vector3d> B1(b1), B2(b2), B3(b3), B4(b4), B5(b5), B6(b6), B7(b7), B8(b8);

    Eigen::Vector3d B11 = B1 / 2.0, B12 = B2 / 2.0, B13 = B3 / 2.0, B14 = B4 / 2.0,
            B15 = B5 / 2.0, B16 = B6 / 2.0, B17 = B7 / 2.0, B18 = B8 / 2.0;

    Eigen::Vector3d B21 = B1 * 2.0, B22 = B2 * 2.0, B23 = B3 * 2.0, B24 = B4 * 2.0,
            B25 = B5 * 2.0, B26 = B6 * 2.0, B27 = B7 * 2.0, B28 = B8 * 2.0;

    points.push_back(B1);
    points.push_back(B2);
    points.push_back(B3);
    points.push_back(B4);
    points.push_back(B5);
    points.push_back(B6);
    points.push_back(B7);
    points.push_back(B8);
    points.push_back(B11);
    points.push_back(B12);
    points.push_back(B13);
    points.push_back(B14);
    points.push_back(B15);
    points.push_back(B16);
    points.push_back(B17);
    points.push_back(B18);
    points.push_back(B21);
    points.push_back(B22);
    points.push_back(B23);
    points.push_back(B24);
    points.push_back(B25);
    points.push_back(B26);
    points.push_back(B27);
    points.push_back(B28);

    g2o::GaussianSampler<Eigen::Vector2d, Eigen::Matrix2d> featNoise;
    featNoise.setDistribution(featVar);

    Eigen::Matrix<double, 6, 6> poseVar = 0.01 * Eigen::Matrix<double, 6, 6>::Identity();
    gpsVar = poseVar;
    g2o::GaussianSampler<Eigen::Matrix<double, 6, 1>, Eigen::Matrix<double, 6, 6>> poseNoise;

    poseNoise.setDistribution(poseVar);

    Eigen::Vector3d zx(0, 0, 1);
    bool plus = true;
    double x = 0.0, y = 0.0, z = 0.0;
    int theta = 45;
    double time = 0.0;
    for(int k = 0; k <= 1000; ++k) {
        time = time + 0.01;
        z = 6.0 - k * 0.012;
        x = r * 1.2 * exp(k / 800.0) * std::cos(3.1415 / 180.0 * theta);
        y = r * exp(k / 1000.0) * std::sin(3.1415 / 180.0 * theta);
        if(theta == 135)
            plus = false;
        if(theta == 45)
            plus = true;
        if(plus)
            theta++;
        else
            theta--;

        Eigen::Vector3d n_(x, y, z);
        auto z1 = -n_.normalized();
        Eigen::Vector3d x1 = -zx.cross(z1);
        Eigen::Vector3d y1 = z1.cross(x1);
        Eigen::Matrix3d R;
        R.block<3, 1>(0, 0) = x1;
        R.block<3, 1>(0, 1) = y1;
        R.block<3, 1>(0, 2) = z1;
        Sophus::SE3d se3(R, n_);
        Pose pose;
        pose.pose = se3;

        pose.projectPoints.push_back(se3 * B1);
        pose.projectPoints.push_back(se3 * B2);
        pose.projectPoints.push_back(se3 * B3);
        pose.projectPoints.push_back(se3 * B4);
        pose.projectPoints.push_back(se3 * B5);
        pose.projectPoints.push_back(se3 * B6);
        pose.projectPoints.push_back(se3 * B7);
        pose.projectPoints.push_back(se3 * B8);
        pose.projectPoints.push_back(se3 * B11);
        pose.projectPoints.push_back(se3 * B12);
        pose.projectPoints.push_back(se3 * B13);
        pose.projectPoints.push_back(se3 * B14);
        pose.projectPoints.push_back(se3 * B15);
        pose.projectPoints.push_back(se3 * B16);
        pose.projectPoints.push_back(se3 * B17);
        pose.projectPoints.push_back(se3 * B18);
        pose.projectPoints.push_back(se3 * B21);
        pose.projectPoints.push_back(se3 * B22);
        pose.projectPoints.push_back(se3 * B23);
        pose.projectPoints.push_back(se3 * B24);
        pose.projectPoints.push_back(se3 * B25);
        pose.projectPoints.push_back(se3 * B26);
        pose.projectPoints.push_back(se3 * B27);
        pose.projectPoints.push_back(se3 * B28);

        pose.noisePose =  Sophus::SE3d::exp(poseNoise.generateSample()) * se3;
        pose.timeStamp = time;
        for(size_t i = 0; i < pose.projectPoints.size(); ++i) {
            Eigen::Vector2d feat_;
            feat_(0) = pose.projectPoints[i](0) / pose.projectPoints[i](2);
            feat_(1) = pose.projectPoints[i](1) / pose.projectPoints[i](2);
            feat_ = feat_ + featNoise.generateSample();
            pose.noiseFeats.push_back(Eigen::Vector2i(int(feat_(0)), int(feat_(1))));
        }

        poses.push_back(pose);
    }

    for(size_t it = 0; it + 1 < poses.size(); ++it) {
        if(it % InterVal) {
            double gtime = get_gpstimestamp(poses[it].timeStamp, poses[it + 1].timeStamp);
            double alpha = (gtime - poses[it].timeStamp) / (poses[it + 1].timeStamp - poses[it].timeStamp);
            auto Tij = poses[it + 1].pose * poses[it].pose.inverse();
            Tij = Sophus::SE3d::exp(alpha * Tij.log()) * poses[it].pose;
            GPSMeasure gps;
            gps.timeStamp = gtime;
            gps.pose = Tij;
            gpsMeasures.push_back(gps);
            gps.pose = Sophus::SE3d::exp(poseNoise.generateSample()) * gps.pose;
            gpsMeasuresWithNoise.push_back(gps);
        }
    }

}

